import os
import random

from moviepy import VideoFileClip, concatenate_videoclips

from preprocess.material_clip import iterative_video_crop, iterative_audio_crop, process_image, \
    split_long_video_to_clips


def collect_media_files(root_dir: str):
    """
    扫描目录及其子目录，将媒体文件按类型分类并返回。

    Returns:
        dict: 包含图片、视频、音频三类文件路径列表
    """
    media_files = {
        'images': [],
        'videos': [],
        'audios': []
    }

    for dirpath, _, filenames in os.walk(root_dir):
        for file in filenames:
            if "剪裁" not in file and "图片处理" not in file:
                file_path = os.path.join(dirpath, file)
                ext = os.path.splitext(file)[1].lower()

                if ext in ('.jpg', '.jpeg', '.png'):
                    media_files['images'].append(file_path)
                elif ext in ('.mp4', '.avi', '.mov'):
                    media_files['videos'].append(file_path)
                elif ext in ('.mp3', '.wav'):
                    media_files['audios'].append(file_path)

    return media_files


def get_random_video_clip(video_path, clip_duration):
    """
    从视频中随机截取指定长度的片段。

    Args:
        video_path (str): 视频文件路径
        clip_duration (float): 截取片段的时长（秒）

    Returns:
        VideoFileClip: 截取的视频片段
    """
    video = VideoFileClip(video_path)
    max_start_time = max(0, video.duration - clip_duration)
    start_time = random.uniform(0, max_start_time)
    end_time = start_time + clip_duration
    return video.subclipped(start_time, end_time)


def mix_videos(video_paths, clip_duration=5):
    """
    混剪多个视频，每个视频随机截取指定长度的片段。

    Args:
        video_paths (list): 视频文件路径列表
        clip_duration (float): 每个片段的时长（秒）

    Returns:
        VideoFileClip: 混剪后的视频
    """
    clips = [get_random_video_clip(path, clip_duration) for path in video_paths]
    return concatenate_videoclips(clips)


def concatenate_short_videos(video_paths, duration=13, clip_duration=5, mix_enabled=True, mix_count=1):
    """
    拼接长度小于 min_duration 的视频，直到总长度大于等于 min_duration，并支持混剪。

    Args:
        video_paths (list): 视频文件路径列表
        duration (float): 最小保留时长（秒）
        clip_duration (float): 每个片段的时长（秒）
        mix_enabled (bool): 是否启用混剪功能
        mix_count (int): 混剪视频的数量

    Returns:
        list: 拼接后的视频文件路径列表
    """
    short_videos = []
    long_videos = []
    concatenated_videos = []

    for video_path in video_paths:
        try:
            clip = VideoFileClip(video_path)
            if clip.duration < duration:
                short_videos.append(video_path)
            else:
                long_videos.append(video_path)
        except Exception as e:
            print(f"[错误1]下一个: {e}")
            continue

    # 处理短视频：混剪
    if mix_enabled and short_videos:
        for _ in range(mix_count):
            current_duration = 0
            current_clips = []

            while current_duration < duration and short_videos:
                try:
                    video_path = random.choice(short_videos)
                    clip = get_random_video_clip(video_path, clip_duration)
                    current_duration += clip.duration
                    current_clips.append(clip)
                except Exception as e:
                    print(f"[错误2],下一个 {output_path}: {e}")
                    continue

            if current_clips:
                try:
                    concatenated_clip = concatenate_videoclips(current_clips)
                    output_path = os.path.join(os.path.dirname(video_paths[0]),
                                               f"混剪_{len(concatenated_videos) + 1}_时长{int(current_duration)}s.mp4")
                    concatenated_clip.write_videofile(output_path, codec="libx264", audio_codec="aac")
                    concatenated_videos.append(output_path)
                except Exception as e:
                    print(f"[错误3] ,下一个 {output_path}: {e}")
                    continue

    # 处理长视频：直接添加
    for video_path in long_videos:
        concatenated_videos.append(video_path)

    return concatenated_videos


def process_media_files(root_dir: str):
    """
    先扫描所有媒体文件，统一处理。
    """
    media_files = collect_media_files(root_dir)

    print(f"共发现 {len(media_files['images'])} 张图片")
    print(f"共发现 {len(media_files['videos'])} 个视频")
    print(f"共发现 {len(media_files['audios'])} 个音频")

    # 处理图片
    for image_path in media_files['images']:
        print(f"[图片] 正在处理: {image_path}")
        try:
            process_image(image_path, )
        except Exception as e:
            print(f"[错误] 图片处理失败 {image_path}: {e}")

    # 处理视频
    if media_files['videos']:
        # 拼接短视频（支持混剪）
        concatenated_videos = concatenate_short_videos(
            media_files['videos'],
            duration=20,
            clip_duration=20,
            mix_enabled=True,  # 启用混剪
            mix_count=20  # 混剪视频数量为 2
        )
        # 画中画：混合模式-不透明度

        # 裁剪拼接后的视频
        for video_path in concatenated_videos:
            print(f"[视频] 正在处理: {video_path}")
            try:
                iterative_video_crop(video_path, min_duration=13, step=9.0)
            except Exception as e:
                print(f"[错误] 视频裁剪失败 {video_path}: {e}")

    # 处理音频
    for audio_path in media_files['audios']:
        print(f"[音频] 正在处理: {audio_path}")
        try:
            iterative_audio_crop(audio_path, min_duration=20, step=9.0)
        except Exception as e:
            print(f"[错误] 音频裁剪失败 {audio_path}: {e}")


def process_videos_to_clips(video_dir, min_duration=20):
    """
    处理多个视频文件，将每个长视频分割为多个最小长度的视频片段。

    Args:
        video_paths (list): 视频文件路径列表
        min_duration (float): 每个片段的最小持续时间（秒）

    Returns:
        list: 所有分割后的视频文件路径列表
    """
    media_files = collect_media_files(video_dir)

    all_clips = []

    if media_files['videos']:
        video_paths = media_files['videos']
        for video_path in video_paths:
            try:
                print(f"[视频] 正在分割: {video_path}")
                clips = split_long_video_to_clips(video_path, min_duration)
                all_clips.extend(clips)
            except Exception as e:
                print(f"[错误] 视频分割失败 {video_path}: {e}")
                all_clips.append(video_path)  # 保留原始视频路径

    return all_clips


if __name__ == "__main__":
    # 设置要处理的根目录
    # base_path="D:\模版9_图片口播"
    # base_path="D:\模版3_穿搭展示"
    # base_path="D:\模版6_稳定性效果"
    # base_path="D:\模版2_产品拼接"
    # base_path="D:\模版5_模特图穿插"
    # base_path="D:\模版1_主题蹭款"
    base_path = "D:\批量_视频保存\测试视频"
    # base_path = "D:\模版4_美甲展示"
    media_root_directory = [
        # os.path.join(base_path, "解压-解绳子"),
        # os.path.join(base_path, "团播-男团"),
        os.path.join(base_path, "扎袋子"),
        # os.path.join(base_path, "图片_2"),
        # os.path.join(base_path, "图片_3"),
        # os.path.join(base_path, "图片_4"),
        # os.path.join(base_path, "图片_5"),
        # os.path.join(base_path, "图片_6"),
        # os.path.join(base_path, "图片_7"),
        # os.path.join(base_path, "图片_8"),
        # os.path.join(base_path, "视频_1"),
        # os.path.join(base_path, "视频_2"),
        # os.path.join(base_path, "音频")
    ]
    for i in media_root_directory:
        process_videos_to_clips(i, min_duration=40)
    # 开始批量处理
    # for i in media_root_directory:
    #
    #     process_media_files(i)
    # process_media_files(media_root_directory)
